package org.xyl.engine;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
//import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

public class ChromaEngine {

    static String myOpenAI_token="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";

    public static void main(String[] args) {
        // 1. 加载文档
        Path documentPath = Paths.get("F:\\架构师材料\\architect-awesome\\README.md"); // 替换为你的文档路径
        DocumentParser documentParser = new TextDocumentParser();
        Document document = FileSystemDocumentLoader.loadDocument(documentPath, documentParser);

        // 2. 分割文档为小块
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 50);
        List<TextSegment> segments = splitter.split(document);

        // 3. 初始化嵌入模型
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        // 4. 为每个文本段生成嵌入向量
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        // 5. 初始化Chroma向量存储
        EmbeddingStore<TextSegment> embeddingStore = ChromaEmbeddingStore.builder()
                .baseUrl("http://localhost:8000") // Chroma服务器地址
                .collectionName("example-collection") // 集合名称
                .build();

        // 6. 将嵌入向量和文本段存储到Chroma
        embeddingStore.addAll(embeddings, segments);

        // 7. 初始化聊天模型
//        ChatLanguageModel chatModel = OpenAiChatModel.builder()
//                .apiKey(myOpenAI_token) // 替换为你的OpenAI API密钥
//                .modelName("gpt-3.5-turbo")
//                .temperature(0.7)
//                        .build();

        // 8. 用户查询
        String userQuery = "什么是人工智能？";

        // 9. 为查询生成嵌入向量
        Embedding queryEmbedding = embeddingModel.embed(userQuery).content();

        // 10. 从Chroma中检索相关文档片段
        int maxResults = 3;
        EmbeddingSearchRequest embeddingSearchRequest = new EmbeddingSearchRequest(queryEmbedding, maxResults, 0.75,null);
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> relevantSegments = search.matches();

        // 11. 构建上下文
        StringBuilder contextBuilder = new StringBuilder();
        for (EmbeddingMatch<TextSegment> match : relevantSegments) {
            contextBuilder.append(match.embedded().text()).append("\n\n");
        }
        String context = contextBuilder.toString();

        // 12. 构建提示
        String prompt = String.format("基于以下上下文回答问题:\n\n%s\n\n问题: %s", context, userQuery);

        // 13. 生成回答
//        String answer = chatModel.generate(prompt);

        // 14. 输出结果
        System.out.println("问题: " + userQuery);
//        System.out.println("回答: " + answer);
    }
}
